#!/usr/bin/env python3
"""
Elasticsearch操作工具类
封装ES的增删改查操作
"""

import asyncio
from typing import List, Dict, Any, Optional
from elasticsearch import Elasticsearch, AsyncElasticsearch
from elasticsearch.exceptions import NotFoundError, RequestError
import json
import logging

from services.recommendation_dsl import RecommendationDSL

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ESClient:
    """Elasticsearch客户端工具类"""
    
    def __init__(self, es_host="http://localhost:9200", use_async=False):
        self.es_host = es_host
        self.use_async = use_async
        
        if use_async:
            self.es = AsyncElasticsearch([es_host])
        else:
            self.es = Elasticsearch([es_host])
        
        self.index_name = "content_recommend"
        self.dsl_builder = RecommendationDSL()
    
    async def health_check(self) -> bool:
        """健康检查"""
        try:
            if self.use_async:
                response = await self.es.ping()
            else:
                response = self.es.ping()
            
            logger.info("ES连接正常")
            return True
        except Exception as e:
            logger.error(f"ES连接失败: {e}")
            return False
    
    async def create_document(self, doc_id: int, document: Dict[str, Any]) -> bool:
        """创建文档"""
        try:
            if self.use_async:
                response = await self.es.index(
                    index=self.index_name,
                    id=doc_id,
                    body=document
                )
            else:
                response = self.es.index(
                    index=self.index_name,
                    id=doc_id,
                    body=document
                )
            
            logger.info(f"文档 {doc_id} 创建成功")
            return True
            
        except Exception as e:
            logger.error(f"创建文档 {doc_id} 失败: {e}")
            return False
    
    async def update_document(self, doc_id: int, document: Dict[str, Any]) -> bool:
        """更新文档"""
        try:
            if self.use_async:
                response = await self.es.index(
                    index=self.index_name,
                    id=doc_id,
                    body=document
                )
            else:
                response = self.es.index(
                    index=self.index_name,
                    id=doc_id,
                    body=document
                )
            
            logger.info(f"文档 {doc_id} 更新成功")
            return True
            
        except Exception as e:
            logger.error(f"更新文档 {doc_id} 失败: {e}")
            return False
    
    async def delete_document(self, doc_id: int) -> bool:
        """删除文档"""
        try:
            if self.use_async:
                response = await self.es.delete(
                    index=self.index_name,
                    id=doc_id
                )
            else:
                response = self.es.delete(
                    index=self.index_name,
                    id=doc_id
                )
            
            logger.info(f"文档 {doc_id} 删除成功")
            return True
            
        except NotFoundError:
            logger.warning(f"文档 {doc_id} 不存在")
            return True
        except Exception as e:
            logger.error(f"删除文档 {doc_id} 失败: {e}")
            return False
    
    async def get_document(self, doc_id: int) -> Optional[Dict[str, Any]]:
        """获取文档"""
        try:
            if self.use_async:
                response = await self.es.get(
                    index=self.index_name,
                    id=doc_id
                )
            else:
                response = self.es.get(
                    index=self.index_name,
                    id=doc_id
                )
            
            return response['_source']
            
        except NotFoundError:
            logger.warning(f"文档 {doc_id} 不存在")
            return None
        except Exception as e:
            logger.error(f"获取文档 {doc_id} 失败: {e}")
            return None
    
    async def search_documents(
        self, 
        query: Dict[str, Any],
        size: int = 20,
        from_: int = 0
    ) -> Dict[str, Any]:
        """搜索文档"""
        try:
            search_body = {
                "query": query,
                "size": size,
                "from": from_
            }
            
            if self.use_async:
                response = await self.es.search(
                    index=self.index_name,
                    body=search_body
                )
            else:
                response = self.es.search(
                    index=self.index_name,
                    body=search_body
                )
            
            return response
            
        except Exception as e:
            logger.error(f"搜索文档失败: {e}")
            return {"hits": {"hits": [], "total": {"value": 0}}}
    
    async def tag_based_recommendation(
        self,
        user_tag_ids: List[int],
        tag_weights: Optional[Dict[int, float]] = None,
        size: int = 20,
        from_: int = 0
    ) -> List[Dict[str, Any]]:
        """基于标签的推荐"""
        try:
            query_dsl = self.dsl_builder.build_tag_matching_query(
                user_tag_ids=user_tag_ids,
                tag_weights=tag_weights,
                size=size,
                from_=from_
            )
            
            response = await self.search_documents(
                query=query_dsl["body"]["query"],
                size=size,
                from_=from_
            )
            
            # 提取推荐结果
            recommendations = []
            for hit in response["hits"]["hits"]:
                doc = hit["_source"]
                doc["_score"] = hit["_score"]
                recommendations.append(doc)
            
            logger.info(f"标签推荐返回 {len(recommendations)} 条结果")
            return recommendations
            
        except Exception as e:
            logger.error(f"标签推荐失败: {e}")
            return []
    
    async def text_similarity_recommendation(
        self,
        content_id: int,
        fields: List[str] = None,
        size: int = 20,
        min_score: float = 0.1
    ) -> List[Dict[str, Any]]:
        """基于文本相似的推荐"""
        try:
            query_dsl = self.dsl_builder.build_text_similarity_query(
                content_id=content_id,
                fields=fields,
                size=size,
                min_score=min_score
            )
            
            response = await self.search_documents(
                query=query_dsl["body"]["query"],
                size=size
            )
            
            # 提取推荐结果
            recommendations = []
            for hit in response["hits"]["hits"]:
                doc = hit["_source"]
                doc["_score"] = hit["_score"]
                recommendations.append(doc)
            
            logger.info(f"文本相似推荐返回 {len(recommendations)} 条结果")
            return recommendations
            
        except Exception as e:
            logger.error(f"文本相似推荐失败: {e}")
            return []
    
    async def hybrid_recommendation(
        self,
        user_tag_ids: List[int],
        reference_content_id: Optional[int] = None,
        tag_weights: Optional[Dict[int, float]] = None,
        size: int = 20,
        from_: int = 0
    ) -> List[Dict[str, Any]]:
        """混合推荐（标签匹配 + 文本相似）"""
        try:
            query_dsl = self.dsl_builder.build_hybrid_recommendation_query(
                user_tag_ids=user_tag_ids,
                reference_content_id=reference_content_id,
                tag_weights=tag_weights,
                size=size,
                from_=from_
            )
            
            response = await self.search_documents(
                query=query_dsl["body"]["query"],
                size=size,
                from_=from_
            )
            
            # 提取推荐结果
            recommendations = []
            for hit in response["hits"]["hits"]:
                doc = hit["_source"]
                doc["_score"] = hit["_score"]
                recommendations.append(doc)
            
            logger.info(f"混合推荐返回 {len(recommendations)} 条结果")
            return recommendations
            
        except Exception as e:
            logger.error(f"混合推荐失败: {e}")
            return []
    
    async def custom_scoring_recommendation(
        self,
        user_tag_ids: List[int],
        tag_weights: Optional[Dict[int, float]] = None,
        boost_factors: Optional[Dict[str, float]] = None,
        size: int = 20,
        from_: int = 0
    ) -> List[Dict[str, Any]]:
        """自定义打分排序推荐"""
        try:
            # 构建基础查询
            base_query = self.dsl_builder.build_tag_matching_query(
                user_tag_ids=user_tag_ids,
                tag_weights=tag_weights,
                size=size,
                from_=from_
            )
            
            # 应用自定义打分
            query_dsl = self.dsl_builder.build_custom_scoring_query(
                base_query=base_query,
                boost_factors=boost_factors
            )
            
            response = await self.search_documents(
                query=query_dsl["body"]["query"],
                size=size,
                from_=from_
            )
            
            # 提取推荐结果
            recommendations = []
            for hit in response["hits"]["hits"]:
                doc = hit["_source"]
                doc["_score"] = hit["_score"]
                recommendations.append(doc)
            
            logger.info(f"自定义打分推荐返回 {len(recommendations)} 条结果")
            return recommendations
            
        except Exception as e:
            logger.error(f"自定义打分推荐失败: {e}")
            return []
    
    async def search_after_recommendation(
        self,
        user_tag_ids: List[int],
        search_after: Optional[List] = None,
        tag_weights: Optional[Dict[int, float]] = None,
        size: int = 20
    ) -> Dict[str, Any]:
        """search_after分页推荐"""
        try:
            query_dsl = self.dsl_builder.build_search_after_query(
                user_tag_ids=user_tag_ids,
                search_after=search_after,
                tag_weights=tag_weights,
                size=size
            )
            
            search_body = query_dsl["body"]
            
            if self.use_async:
                response = await self.es.search(
                    index=self.index_name,
                    body=search_body
                )
            else:
                response = self.es.search(
                    index=self.index_name,
                    body=search_body
                )
            
            # 提取推荐结果
            recommendations = []
            for hit in response["hits"]["hits"]:
                doc = hit["_source"]
                doc["_score"] = hit["_score"]
                recommendations.append(doc)
            
            # 获取下一页的search_after值
            next_search_after = None
            if response["hits"]["hits"]:
                next_search_after = response["hits"]["hits"][-1]["sort"]
            
            result = {
                "recommendations": recommendations,
                "next_search_after": next_search_after,
                "total": response["hits"]["total"]["value"]
            }
            
            logger.info(f"search_after推荐返回 {len(recommendations)} 条结果")
            return result
            
        except Exception as e:
            logger.error(f"search_after推荐失败: {e}")
            return {
                "recommendations": [],
                "next_search_after": None,
                "total": 0
            }
    
    async def bulk_index(self, documents: List[Dict[str, Any]]) -> bool:
        """批量索引文档"""
        try:
            bulk_data = []
            for doc in documents:
                bulk_data.append({
                    "index": {
                        "_index": self.index_name,
                        "_id": doc["content_id"]
                    }
                })
                bulk_data.append(doc)
            
            if self.use_async:
                response = await self.es.bulk(body=bulk_data)
            else:
                response = self.es.bulk(body=bulk_data)
            
            logger.info(f"批量索引 {len(documents)} 个文档成功")
            return True
            
        except Exception as e:
            logger.error(f"批量索引失败: {e}")
            return False
    
    async def close(self):
        """关闭连接"""
        if self.use_async:
            await self.es.close()

# 同步版本的ES客户端
class SyncESClient(ESClient):
    """同步版本的ES客户端"""
    
    def __init__(self, es_host="localhost:9200"):
        super().__init__(es_host, use_async=False)
    
    def tag_based_recommendation_sync(
        self,
        user_tag_ids: List[int],
        tag_weights: Optional[Dict[int, float]] = None,
        size: int = 20,
        from_: int = 0
    ) -> List[Dict[str, Any]]:
        """同步版本的标签推荐"""
        return asyncio.run(
            self.tag_based_recommendation(
                user_tag_ids, tag_weights, size, from_
            )
        )
    
    def text_similarity_recommendation_sync(
        self,
        content_id: int,
        fields: List[str] = None,
        size: int = 20,
        min_score: float = 0.1
    ) -> List[Dict[str, Any]]:
        """同步版本的文本相似推荐"""
        return asyncio.run(
            self.text_similarity_recommendation(
                content_id, fields, size, min_score
            )
        )

if __name__ == "__main__":
    # 测试ES客户端
    async def test_es_client():
        es_client = ESClient(use_async=True)
        
        # 健康检查
        health = await es_client.health_check()
        print(f"ES健康状态: {health}")
        
        # 测试标签推荐
        recommendations = await es_client.tag_based_recommendation(
            user_tag_ids=[1, 3, 5],
            tag_weights={1: 2.0, 3: 1.5, 5: 1.0},
            size=10
        )
        
        print(f"推荐结果数量: {len(recommendations)}")
        
        await es_client.close()
    
    asyncio.run(test_es_client())
